package com.emeraldparser.util;

import java.util.Collection;
import java.util.Iterator;
import java.util.Set;

import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.ForwardingIterator;
import com.google.common.collect.ForwardingSet;
import com.google.common.collect.Multimap;

/**
 * A forwarding set that indexes its elements by a key that is computed for each element.
 * 
 * @author Adam Paynter
 *
 * @param <E>
 * @param <K>
 */
public class SurjectiveMappingSet<E, K> extends ForwardingSet<E> {

	private final Set<E> delegate;
	private final Function<? super E, ? extends K> function;
	private final Multimap<? super K, ? super E> multimap;

	public SurjectiveMappingSet(Set<E> delegate, Function<? super E, ? extends K> function, Multimap<? super K, ? super E> multimap) {
		this.delegate = Preconditions.checkNotNull(delegate, "Set must not be null");
		this.function = Preconditions.checkNotNull(function, "Function must not be null");
		this.multimap = Preconditions.checkNotNull(multimap, "Multimap must not be null");
	}

	@Override
	protected Set<E> delegate() {
		return delegate;
	}

	private K key(E element) {
		return function.apply(element);
	}

	@Override
	public boolean add(E element) {
		boolean modified = super.add(element);
		
		if (modified) {
			multimap.put(key(element), element);
		}
		
		return modified;
	}

	@Override
	public boolean addAll(Collection<? extends E> collection) {
		return standardAddAll(collection);
	}

	@Override
	public void clear() {
		standardClear();
	}

	@Override
	public Iterator<E> iterator() {
		final Iterator<E> iterator = super.iterator();
		
		return new ForwardingIterator<E>() {

			private E current;
			
			@Override
			protected Iterator<E> delegate() {
				return iterator;
			}

			@Override
			public E next() {
				return current = super.next();
			}

			@Override
			public void remove() {
				super.remove();
				elementRemoved(current);
			}
			
		};
	}

	@Override
	public boolean remove(Object object) {
		boolean modified = super.remove(object);
		
		if (modified) {
			@SuppressWarnings("unchecked")
			E element = (E) object; // Because object was removed from the set, we know that it must be of type E
			elementRemoved(element);
		}
		
		return modified;
	}

	private void elementRemoved(E element) {
		multimap.remove(key(element), element);
	}

	@Override
	public boolean removeAll(Collection<?> collection) {
		return standardRemoveAll(collection);
	}

	@Override
	public boolean retainAll(Collection<?> collection) {
		return standardRetainAll(collection);
	}

}
